#!/bin/bash

# usage:
# in aistudio code server terminal
# curl https://gitee.com/liuwake/Cloud/raw/master/Baidu_AIStudio/bd_mmpre.sh | bash

# add in jupyter-notebook cell
# import os ; os.environ["PATH"] += os.pathsep + os.path.expanduser("~/.local/bin")
# %cd /home/aistudio/mmpretrain
# !mim train mmpre configs/resnet/resnet18_8xb16_cifar10.py \
#     --amp

build_mmpre() {
    export PATH=/home/aistudio/.local/bin:$PATH
    echo $PATH

    python -m site

    wget -O /home/aistudio/btop https://r2.neng.buzz/app/btop.conf
    wget -O /home/aistudio/btop https://r2.neng.buzz/app/btop \
        && chmod +x /home/aistudio/btop
    wget https://r2.neng.buzz/app/iliazeus.vscode-ansi-1.1.7.vsix \
        -O /home/aistudio/iliazeus.vscode-ansi-1.1.7.vsix

    pip install --user \
        openmim numpy==1.26.4 \
        torch==1.11.0+cu113 torchvision==0.12.0+cu113 \
        -f https://mirrors.aliyun.com/pytorch-wheels/cu113/

    mim install -q --user \
        mmcv==2.1.0 mmpretrain
    python -c "import mmpretrain; print(mmpretrain.__version__)"

    cp -r /home/aistudio/.local /home/aistudio/local

    git clone https://gitee.com/liuwake/mmpretrain.git /home/aistudio/mmpretrain

    mkdir -p /home/aistudio/mmpretrain/data/cifar10 \
        && [ ! -e /home/aistudio/mmpretrain/data/cifar10/cifar-10-python.tar.gz ] \
        && wget -O /home/aistudio/mmpretrain/data/cifar10/cifar-10-python.tar.gz \
        https://r2.neng.buzz/data/cifar-10-python.tar.gz
}

set_mmpre() {

    ln -sf /home/aistudio/local/bin /home/aistudio/.local/
    ln -sf /home/aistudio/local/lib /home/aistudio/.local/
    mkdir -p /home/aistudio/.config/btop
    ln -sf /home/aistudio/btop.conf /home/aistudio/.config/btop/btop.conf
    /opt/conda/include/vscode/vscode/bin/remote-cli/code --install-extension /home/aistudio/iliazeus.vscode-ansi-1.1.7.vsix
}

if [ ! -d "/home/aistudio/mmpretrain" ] || [ ! -d "/home/aistudio/local" ]; then
    build_mmpre
else
    set_mmpre
fi
